import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button
import matplotlib.patches as patches

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class TwoLinkArm:
    def __init__(self, l1_range=(100, 150), l2=50):
        """
        初始化双连杆机械臂
        :param l1_range: L1的长度范围 (min, max)
        :param l2: L2的长度
        """
        self.l1_min, self.l1_max = l1_range
        self.l1 = (l1_range[0] + l1_range[1]) / 2  # 初始长度为范围中点
        self.l2 = l2
        
        # 初始角度 (弧度)
        self.theta1 = 0
        self.theta2 = 0
        
        # 创建图形和轴
        self.fig, self.ax = plt.subplots(figsize=(10, 8))
        plt.subplots_adjust(left=0.1, bottom=0.4)  # 留出底部空间放置滑块
        
        # 设置坐标轴范围
        max_reach = self.l1_max + self.l2
        self.ax.set_xlim(-max_reach * 1.1, max_reach * 1.1)
        self.ax.set_ylim(-max_reach * 1.1, max_reach * 1.1)
        self.ax.set_aspect('equal')
        self.ax.grid(True)
        self.ax.set_title('双连杆机械臂运动学分析')
        self.ax.set_xlabel('X坐标')
        self.ax.set_ylabel('Y坐标')
        
        # 绘制底座
        self.base = plt.Circle((0, 0), 10, color='gray')
        self.ax.add_patch(self.base)
        
        # 初始化连杆和末端执行器
        self.link1, = self.ax.plot([], [], 'b-', lw=6)
        self.link2, = self.ax.plot([], [], 'r-', lw=6)
        self.end_effector, = self.ax.plot([], [], 'go', markersize=10)
        
        # 绘制工作空间范围
        self.max_circle = plt.Circle((0, 0), self.l1_max + self.l2, 
                                    color='lightgray', fill=True, alpha=0.3)
        self.min_circle = plt.Circle((0, 0), abs(self.l1_min - self.l2), 
                                    color='white', fill=True)
        self.ax.add_patch(self.max_circle)
        self.ax.add_patch(self.min_circle)
        
        # 创建滑块
        self._create_sliders()
        
        # 更新显示
        self.update_arm()
        
        # 计算并显示工作空间采样点
        self.plot_workspace()
        
    def _create_sliders(self):
        """创建交互滑块控件"""
        # L1长度滑块
        ax_l1 = plt.axes([0.1, 0.25, 0.8, 0.03])
        self.slider_l1 = Slider(ax_l1, 'L1长度', 
                              self.l1_min, self.l1_max, 
                              valinit=self.l1)
        
        # θ1角度滑块 (度)
        ax_theta1 = plt.axes([0.1, 0.2, 0.8, 0.03])
        self.slider_theta1 = Slider(ax_theta1, 'θ1角度 (°)', 
                                  0, 360, valinit=self.theta1 * 180 / np.pi)
        
        # θ2角度滑块 (度)
        ax_theta2 = plt.axes([0.1, 0.15, 0.8, 0.03])
        self.slider_theta2 = Slider(ax_theta2, 'θ2角度 (°)', 
                                  -180, 180, valinit=self.theta2 * 180 / np.pi)
        
        # L2长度滑块
        ax_l2 = plt.axes([0.1, 0.1, 0.8, 0.03])
        self.slider_l2 = Slider(ax_l2, 'L2长度', 
                              10, 100, valinit=self.l2)
        
        # 重置按钮
        ax_reset = plt.axes([0.4, 0.02, 0.2, 0.04])
        self.button_reset = Button(ax_reset, '重置')
        
        # 绑定事件处理函数
        self.slider_l1.on_changed(self.update_l1)
        self.slider_theta1.on_changed(self.update_theta1)
        self.slider_theta2.on_changed(self.update_theta2)
        self.slider_l2.on_changed(self.update_l2)
        self.button_reset.on_clicked(self.reset)
        
    def calculate_position(self):
        """计算连杆端点位置"""
        # 转换角度为弧度
        theta1_rad = np.radians(self.theta1)
        theta2_rad = np.radians(self.theta2)
        
        # 计算关节点坐标
        x1 = self.l1 * np.cos(theta1_rad)
        y1 = self.l1 * np.sin(theta1_rad)
        
        # 计算末端执行器坐标
        x2 = x1 + self.l2 * np.cos(theta1_rad + theta2_rad)
        y2 = y1 + self.l2 * np.sin(theta1_rad + theta2_rad)
        
        return (x1, y1), (x2, y2)
    
    def update_arm(self):
        """更新机械臂的显示"""
        (x1, y1), (x2, y2) = self.calculate_position()
        
        # 更新连杆位置
        self.link1.set_data([0, x1], [0, y1])
        self.link2.set_data([x1, x2], [y1, y2])
        self.end_effector.set_data([x2], [y2])
        
        # 更新工作空间显示
        max_reach = self.l1 + self.l2
        min_reach = abs(self.l1 - self.l2)
        self.max_circle.set_radius(max_reach)
        self.min_circle.set_radius(min_reach)
        
        # 添加末端位置文本
        self.ax.set_title(f'双连杆机械臂运动学分析 - 末端位置: ({x2:.1f}, {y2:.1f})')
        
        self.fig.canvas.draw_idle()
    
    def plot_workspace(self):
        """绘制工作空间采样点"""
        # 采样角度
        thetas1 = np.linspace(0, 2*np.pi, 100)
        thetas2 = np.linspace(0, 2*np.pi, 100)
        
        # 计算所有可能的末端位置
        workspace_points = []
        for theta1 in thetas1:
            for theta2 in thetas2:
                x = self.l1 * np.cos(theta1) + self.l2 * np.cos(theta1 + theta2)
                y = self.l1 * np.sin(theta1) + self.l2 * np.sin(theta1 + theta2)
                workspace_points.append([x, y])
        
        workspace_points = np.array(workspace_points)
        self.workspace_plot = self.ax.scatter(
            workspace_points[:, 0], workspace_points[:, 1], 
            c='green', s=1, alpha=0.2, label='工作空间'
        )
        self.ax.legend()
    
    def update_workspace(self):
        """更新工作空间采样点"""
        # 先移除旧的工作空间点
        if hasattr(self, 'workspace_plot'):
            self.workspace_plot.remove()
        
        # 绘制新的工作空间点
        self.plot_workspace()
    
    # 滑块更新函数
    def update_l1(self, val):
        self.l1 = val
        self.update_arm()
        self.update_workspace()
    
    def update_theta1(self, val):
        self.theta1 = val
        self.update_arm()
    
    def update_theta2(self, val):
        self.theta2 = val
        self.update_arm()
    
    def update_l2(self, val):
        self.l2 = val
        self.update_arm()
        self.update_workspace()
    
    def reset(self, event):
        """重置所有参数到初始值"""
        self.slider_l1.reset()
        self.slider_theta1.reset()
        self.slider_theta2.reset()
        self.slider_l2.reset()
        self.update_workspace()

if __name__ == "__main__":
    # 创建机械臂实例，L1范围100-150，L2初始值50
    arm = TwoLinkArm(l1_range=(50, 150), l2=50)
    plt.show()
